from flask import Flask, jsonify, request
import threading
import time
import base64
import socket
import pickle
import struct
from io import BytesIO

# === 设置 Matplotlib 使用非 GUI 后端 ===
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

app = Flask(__name__)
SERVER_ADDRESS = ("127.0.0.1", 9999)
SizeStruct = struct.Struct("!I")

class SocketManager:
    def __init__(self, address):
        self.address = address

    def __enter__(self):
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect(self.address)
        return self.sock

    def __exit__(self, *ignore):
        self.sock.close()

# === 后台通信模块 ===
def send_request(command, *args):
    try:
        with SocketManager(SERVER_ADDRESS) as s:
            payload = pickle.dumps((command, *args), 3)
            s.sendall(SizeStruct.pack(len(payload)))
            s.sendall(payload)

            size_data = b""
            while len(size_data) < 4:
                more = s.recv(4 - len(size_data))
                if not more:
                    raise ConnectionError("Connection closed before receiving response size")
                size_data += more

            size = SizeStruct.unpack(size_data)[0]
            data = b""
            while len(data) < size:
                more = s.recv(size - len(data))
                if not more:
                    break
                data += more

            return pickle.loads(data)
    except Exception as e:
        return False, str(e)

# === 后台任务线程 ===
def run_background_socket():
    while True:
        time.sleep(10)

# === 数据分析模块 ===
class DataAnalyze:
    def __init__(self, usage_file):
        self.usage_file = usage_file

    def DrawMonth(self, user_data):
        dates = [entry[1]['date'] for entry in user_data['daily']]
        total_times = [entry[1].get('totaltime', entry[1].get('Total')) for entry in user_data['daily']]
        video_times = [
            entry[1].get('videotime',
                         entry[1].get('TikTok', 0) +
                         entry[1].get('Youtube', 0) +
                         entry[1].get('Instagram', 0))
            for entry in user_data['daily']
        ]
        fig, ax = plt.subplots(figsize=(14, 6))
        ax.plot(dates, total_times, marker='o', linestyle='-', color='b', label='Total Time')
        ax.plot(dates, video_times, marker='s', linestyle='--', color='r', label='Video Time')
        ax.set_title('Daily Total Time and Video Time (Weekly Highlighted with Divider)')
        ax.set_xlabel('Date')
        ax.set_ylabel('Time (hours)')
        ax.grid(True, linestyle='--', alpha=0.5)
        ax.legend()
        ax.set_xticks(range(len(dates)))
        ax.set_xticklabels(dates, rotation=45)
        plt.tight_layout()
        img = BytesIO()
        plt.savefig(img, format='png', dpi=100)
        plt.close()
        return base64.b64encode(img.getvalue()).decode('utf-8')

    def DrawWeeks(self, user_daily_data, days_per_week=7):
        full_dates = [entry[1]['date'] for entry in user_daily_data]
        weeks_data = {}
        for idx, (date, entry) in enumerate(user_daily_data):
            week_num = idx // days_per_week + 1
            if week_num not in weeks_data:
                weeks_data[week_num] = []
            weeks_data[week_num].append((date, entry))

        images = []
        for week_num, week_records in weeks_data.items():
            week_dates = [record[1]['date'] for record in week_records]
            start_date = week_dates[0]
            start_index = full_dates.index(start_date)
            expected_week_dates = full_dates[start_index:start_index + days_per_week]

            total_times = []
            video_times = []
            plotted_dates = []

            for date in expected_week_dates:
                if date in week_dates:
                    idx_in_week = week_dates.index(date)
                    entry = week_records[idx_in_week][1]
                    total_time = entry.get('totaltime', entry.get('Total'))
                    if 'videotime' in entry:
                        video_time = entry['videotime']
                    else:
                        video_time = entry.get('TikTok', 0) + entry.get('Youtube', 0) + entry.get('Instagram', 0)
                else:
                    total_time = 0
                    video_time = 0
                plotted_dates.append(date)
                total_times.append(total_time)
                video_times.append(video_time)

            plt.figure(figsize=(8, 4))
            plt.plot(plotted_dates, total_times, marker='o', linestyle='-', color='b', label='Total Time')
            plt.plot(plotted_dates, video_times, marker='s', linestyle='--', color='r', label='Video Time')
            plt.title(f"Week {week_num}")
            plt.xlabel("Date")
            plt.ylabel("Time (hours)")
            plt.grid(True, linestyle='--', alpha=0.5)
            plt.xticks(rotation=45)
            plt.tight_layout()
            plt.legend()

            img = BytesIO()
            plt.savefig(img, format='png', dpi=100)
            plt.close()
            images.append(base64.b64encode(img.getvalue()).decode('utf-8'))

        return images

    def DrawDay(self, data):
        date_key = data[1]['date']
        keys = [key for key in data[1].keys() if key not in ['date', 'Total']]
        values = [data[1][key] for key in keys]
        total_value = data[1]['Total']

        # 确定纵坐标的上限
        max_value = max(values + [total_value])
        ylim_top = max_value + 1  # 上限设置为最大值加1

        fig, ax = plt.subplots(figsize=(10, 6))  # 增加图表尺寸
        bars = ax.bar(keys, values, color='skyblue', label='Individual Time')
        ax.axhline(y=total_value, color='r', linestyle='--', label='Total Time')

        # 设置纵坐标的范围
        ax.set_ylim(0, ylim_top)

        for bar in bars:
            yval = bar.get_height()
            ax.text(bar.get_x() + bar.get_width() / 2.0, yval + 0.05, round(yval, 2), ha='center', va='bottom')

        ax.set_title(f"Time Spent on Different Platforms on {date_key}")
        ax.set_xlabel("Platform")
        ax.set_ylabel("Time (hours)")
        ax.legend()

        plt.tight_layout()

        img = BytesIO()
        plt.savefig(img, format='png', dpi=100)
        plt.close()
        return base64.b64encode(img.getvalue()).decode('utf-8')


# === Flask HTTP 接口 ===
@app.route('/login', methods=['POST'])
def login_user():
    try:
        data = request.get_json()
        username = data.get('username')
        password = data.get('password')
        ok, msg = send_request("LOGIN", username, password)
        return jsonify({'success': ok, 'message': msg})
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)}), 500

@app.route('/trade/create', methods=['POST'])
def create_trade():
    data = request.get_json()
    ok, msg = send_request("CREATE_TRADE", data)
    return jsonify({"success": ok, "message": msg})

@app.route('/trades', methods=['GET'])
def get_trade():
    data = send_request("GET_TRADES")
    return jsonify(data)

@app.route('/api/analysis/<username>', methods=['GET'])
def api_analysis(username):
    all_data = send_request("GET_USER_DATA")
    if not isinstance(all_data, dict) or username not in all_data:
        return jsonify({"error": "User not found"})

    analyzer = DataAnalyze(all_data)
    user_data = all_data[username]

    # 生成图表并确保返回的是有效的 Base64 编码图片数据
    month_chart = analyzer.DrawMonth(user_data)
    week_charts = analyzer.DrawWeeks(user_data['daily'])
    day_chart = analyzer.DrawDay(user_data['daily'][29])
    response_data = {
        "username": username,
        "user_data": user_data,
        "charts": {
            "month_chart": month_chart,
            "week_charts": week_charts,
            "day_chart": day_chart
        }
    }


    return jsonify(response_data)

@app.route('/post', methods=['POST'])
def create_post():
    data = request.get_json()
    send_request("CREATE_POST", data)
    return jsonify({"status": "success"}), 200

@app.route('/posts', methods=['GET'])
def get_posts():
    data = send_request("GET_POSTS")
    return jsonify(data)

@app.route('/token/<username>', methods=['GET'])
def get_token(username):
    data = send_request("GET_TOKEN",username)
    return jsonify(data)
@app.route('/set_token', methods=['POST'])
def set_token():
    data = request.get_json()
    return jsonify(send_request("SET_TOKEN",data))
# === 启动 ===
if __name__ == '__main__':
    socket_thread = threading.Thread(target=run_background_socket, daemon=True)
    socket_thread.start()
    app.run(host='0.0.0.0', port=5000, debug=True)
